import json
import os
os.environ['TORCH_CUDA_ARCH_LIST'] = '8.6'
from docling.document_converter import DocumentConverter, ImageFormatOption, PowerpointFormatOption, WordFormatOption, \
    ExcelFormatOption, HTMLFormatOption
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions, EasyOcrOptions, AcceleratorDevice
from docling.document_converter import PdfFormatOption, DocumentConverter

# 配置OCR模型，设置EasyOCR模型的路径
easyocr_model_storage_directory = "C:\\Users\\25025\\.cache\\modelscope\\hub\\models\\Ceceliachenen\\easyocr"
easyocr_options = EasyOcrOptions()
easyocr_options.model_storage_directory = easyocr_model_storage_directory

# 配置pdf模型，设置Docling模型的路径
pdf_artifacts_path = "C:\\Users\\25025\\.cache\\modelscope\\hub\\models\\AI-ModelScope\\docling-models"
pdf_pipeline_options = PdfPipelineOptions(artifacts_path=pdf_artifacts_path)
pdf_pipeline_options.ocr_options = easyocr_options
# 设置 accelerator_options.device 为 CUDA 以启用 GPU 加速
pdf_pipeline_options.accelerator_options.device = AcceleratorDevice.CUDA

# 转换模型
converter = DocumentConverter(
    format_options={
        InputFormat.PDF: PdfFormatOption(pipeline_options=pdf_pipeline_options),
        InputFormat.IMAGE: ImageFormatOption(pipeline_options=pdf_pipeline_options),
        InputFormat.PPTX: PowerpointFormatOption(pipeline_options=pdf_pipeline_options),
        InputFormat.DOCX: WordFormatOption(pipeline_options=pdf_pipeline_options),
        InputFormat.XLSX: ExcelFormatOption(pipeline_options=pdf_pipeline_options),
        InputFormat.HTML: HTMLFormatOption(pipeline_options=pdf_pipeline_options)
    }
)

if __name__ == "__main__":

    # 1.1转markdown
    # 定义输入目录和输出目录（请根据实际情况修改路径）
    source_dir = r"C:\Desktop\ZKYU\知识库文件\空气质量标准相关"
    output_dir = r"C:\Desktop\2md\空气质量标准相关"

    # 如果输出目录不存在，则创建
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    # 遍历输入目录下的所有文件
    for filename in os.listdir(source_dir):
        file_path = os.path.join(source_dir, filename)
        if os.path.isfile(file_path):
            print(filename)
            result = converter.convert(file_path)
            markdown_content = result.document.export_to_markdown()
            # 构造输出文件路径，保持原文件名，只更改扩展名为 .md
            output_filename = os.path.splitext(filename)[0] + ".md"
            output_path = os.path.join(output_dir, output_filename)

            # 将 Markdown 内容写入输出文件
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(markdown_content)

            print(f"已转换: {filename} -> {output_filename}")
        else:
            continue

# # 1.2分块
# chunks = list(HierarchicalChunker().chunk(result.document))
#
#
# # 重新定义 `chunk_to_dict`
# def chunk_to_dict(chunk):
#     return {
#         'text': chunk.text,  # 主要文本内容
#         'metadata': {
#             'filename': chunk.meta.origin.filename,  # 获取 PDF 文件名
#             'mimetype': chunk.meta.origin.mimetype,  # 获取文件类型
#             'version': chunk.meta.version,  # 文档版本
#             'doc_items': [
#                 item.text if hasattr(item, 'text') else str(item)  # 兼容 TableItem 和其他类型
#                 for item in chunk.meta.doc_items
#             ]
#         }
#     }
#
#
#
# # 序列化 JSON 并保存
# serializable_chunks = [chunk_to_dict(chunk) for chunk in chunks]
#
# with open("chunks.json", "w", encoding="utf-8") as f:
#     json.dump(serializable_chunks, f, ensure_ascii=False, indent=4)
